{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1572bd9a",
   "metadata": {},
   "source": [
    "# An example generating and displaying the segmentation results\n",
    "This example notebook shows how you can evaluate a pre-trained BRepNet model and display the segmentation results.  We provide some example data and a pretrained model for the evaluation.\n",
    "\n",
    "This notebook should be started from the main BRepNet folder.  Below we check we are in the correct dir and change to the root of the BRepNet repo if not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "436f36e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "if os.path.isfile('../models/brepnet.py'):\n",
    "    # We are in the notebooks directory.  Change to the root\n",
    "    os.chdir('../')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b89b1507",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "# This code allows you to evaluate a pre-trained model for all step files in a folder\n",
    "from eval.evaluate_folder import evaluate_folder\n",
    "\n",
    "# This viewer allows you to visualize the results\n",
    "from visualization.jupyter_segmentation_viewer import JupyterSegmentationViewer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c87e3d6",
   "metadata": {},
   "source": [
    "Now we would like to evaluate the model for some example files.  We can do this using `eval/evaluate_folder.py`.  We need to supply the script with the path to the step files to evaluate, the feature standadization and the pretrained model to use.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cb69622e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 25/25 [01:51<00:00,  4.46s/it]\n",
      "GPU available: True, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed pipeline/extract_feature_data_from_step.py\n",
      "Using labels from example_files/step_examples/temp_working\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/lambouj/anaconda3/envs/brepnet/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py:1585: UserWarning: GPU available but not used. Set the gpus flag in your trainer `Trainer(gpus=1)` or script `--gpus=1`.\n",
      "  \"GPU available but not used. Set the gpus flag in your trainer `Trainer(gpus=1)` or script `--gpus=1`.\"\n",
      "/home/lambouj/anaconda3/envs/brepnet/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py:133: UserWarning: The dataloader, test_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 36 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  f\"The dataloader, {name}, does not have many workers which may be a bottleneck.\"\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7b33a6645f51441aaafbae72958bd5f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "DATALOADER:0 TEST RESULTS\n",
      "{'test/Chamfer_iou': 0.8418079018592834,\n",
      " 'test/CutEnd_iou': 0.7290322780609131,\n",
      " 'test/CutSide_iou': 0.8102856874465942,\n",
      " 'test/ExtrudeEnd_iou': 0.7076271176338196,\n",
      " 'test/ExtrudeSide_iou': 0.8099502325057983,\n",
      " 'test/Fillet_iou': 0.9281437397003174,\n",
      " 'test/RevolveEnd_iou': 0.8999999761581421,\n",
      " 'test/RevolveSide_iou': 0.7386363744735718,\n",
      " 'test/accuracy': 0.8996027708053589,\n",
      " 'test/mean_iou': 0.8081854581832886}\n",
      "--------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# Here is the path to some example step files for us to convert\n",
    "step_folder = Path(\"./example_files/step_examples\")\n",
    "\n",
    "# We will also need to know the feature standardization for the dataset used to train the model\n",
    "# This is found in the dataset file created by pipeline/build_dataset_file.py or pipeline/quickstart.py\n",
    "feature_standardization = Path(\"./example_files/feature_standardization/s2.0.0_step_all_features.json\")\n",
    "\n",
    "# Here is the path to a pretrained model\n",
    "pretrained_model = Path(\"./example_files/pretrained_models/pretrained_s2.0.0_extended_step_uv_net_features_0816_183419.ckpt\")\n",
    "\n",
    "# Now we can evaluate the model on these step files.\n",
    "# Depending on your system you may see pytorch lightning warning you  \n",
    "# GPUs are unused and more worker threads could be used in the dataloader.\n",
    "# The default options here are intended to work on a minimal system.\n",
    "evaluate_folder(step_folder, feature_standardization, pretrained_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c27a8a70",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We found 25 example files\n",
      "Viewing example 142473_f37f7cba_0\n"
     ]
    }
   ],
   "source": [
    "logits_folder = Path(\"example_files/step_examples/temp_working/logits\")\n",
    "step_file_stems = [ f.stem for f in step_folder.glob(\"*.stp\")]\n",
    "print(f\"We found {len(step_file_stems)} example files\")\n",
    "\n",
    "example_index = 0\n",
    "file_stem = step_file_stems[example_index]\n",
    "print(f\"Viewing example {file_stem}\")\n",
    "viewer = JupyterSegmentationViewer(file_stem, step_folder, seg_folder=step_folder, logit_folder=logits_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd35e060",
   "metadata": {},
   "source": [
    "The solid for looks like this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "407a0d40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5bb6b5b2123b4ebeb3a67dd755799ba0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "viewer.view_solid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18e14778",
   "metadata": {},
   "source": [
    "The ground truth segmentation  looks like this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0d85f794",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c3e9c64b06942e7a7dd5b3f3836a5b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "viewer.view_segmentation()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd505370",
   "metadata": {},
   "source": [
    "The predicted segmentation looks like this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fb95f0c2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6bfd55d377f43f588cc02bebacc7518",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "viewer.view_predicted_segmentation()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86fddba1",
   "metadata": {},
   "source": [
    "Here we show the errors in the predicted segmentation.  Faces shown in green are correct and faces shown in red are incorrect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d28594ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fe523e7fd8a54fa69ce2db6ae10ad03f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "viewer.view_errors_in_segmentation()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
